#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import os
import argparse
import pathlib

import numpy as np
import cv2 as cv
        

progress_bar_ratio=0

class PVD(object):

    """PVD image steganographic method"""

    def __init__(self, rangeBits=[3, 3, 4, 5, 6, 7]):
        self.rangeBits = rangeBits
        self.rangeWidth = [np.power(2, x) for x in rangeBits]
        self.ranges = [sum(self.rangeWidth[0:i+1]) for i in range(len(self.rangeWidth))]
        self.ranges.insert(0, 0)
        
    def dInWhichRange(self, d):
        """return the range index of (lk, uk) for d

        """
        for i in range(len(self.ranges)-1):
            if d >= self.ranges[i] and d < self.ranges[i+1]:
                return i
            
    def computeNewD(self, gi, gii, secretBits):
        gi = int(gi)
        gii = int(gii)
        d = gii - gi
        i = self.dInWhichRange(abs(d))

        checkPair = self.pvdFunc(gi, gii, self.ranges[i+1]-d)
        if 0<checkPair[0]<255 and 0<checkPair[1]<255:
            secretValue = np.packbits(secretBits[0:self.rangeBits[i]][::-1], bitorder='little')[0]
            if d >= 0:
                dn = self.ranges[i] + secretValue
            else:
                dn = -(self.ranges[i]+secretValue)
            m = dn - d
            ret = self.pvdFunc(gi, gii, m)
            ret.append(self.rangeBits[i])
            return ret
        else:
            return None

    def extractFromPair(self, gi, gii):
        gi = int(gi)
        gii = int(gii)
        d = gii - gi
        i = self.dInWhichRange(abs(d))

        checkPair = self.pvdFunc(gi, gii, self.ranges[i+1]-d)
        if 0<checkPair[0]<255 and 0<checkPair[1]<255:
            secretValue = abs(d) - self.ranges[i]
            secretBits = np.unpackbits(np.array([secretValue], dtype=np.uint8))[-self.rangeBits[i]:]
            return secretBits
        else:
            return None
    
    def pvdFunc(self, gi, gii, m):
        d = gii - gi
        if d % 2 == 1:
            return [gi-np.ceil(m/2), gii+np.floor(m/2)]
        else:
            return [gi-np.floor(m/2), gii+np.ceil(m/2)]

    def zigzagNextIndex(self, imgRowFlag, imgColFlag, rowlen):
        if imgRowFlag%2 == 0:
            # Even row. Left to right
            if imgColFlag+1 == rowlen:
                # to the end
                newRowFlag = imgRowFlag+1
                newColFlag = imgColFlag
            else:
                newRowFlag = imgRowFlag
                newColFlag = imgColFlag+1
        else:
            # Odd row. Right to left
            if imgColFlag == 0:
                # to the end
                newRowFlag = imgRowFlag+1
                newColFlag = imgColFlag
            else:
                newRowFlag = imgRowFlag
                newColFlag = imgColFlag-1
        return [newRowFlag, newColFlag]

    def hide(self, coverImgPath, secretMsg):
        msg = secretMsg
        msg = np.unpackbits(np.frombuffer(bytes(msg, "utf8"), dtype="uint8"))
        endFlag = np.array([0 for i in range(32)])
        msg = np.append(msg, endFlag)
        count = msg.shape[0]
        io=cv.imread(coverImgPath, cv.IMREAD_GRAYSCALE) # 读入载体图像
        iw=np.array(io)
        # cv.imshow("gray", io)
        # cv.waitKey(0)
        secretFlag = 0
        imgRowFlag = 0
        imgColFlag = 0

        while secretFlag < count:
            if imgRowFlag == io.shape[0]:
                print("Secret message too long. Can not be hidden in cover image")
                break
            (newRowFlag, newColFlag) =  self.zigzagNextIndex(imgRowFlag, imgColFlag, io.shape[1])
            gi = io[imgRowFlag][imgColFlag]
            gii = io[newRowFlag][newColFlag]
            ret = self.computeNewD(gi, gii, msg[secretFlag:])
            if ret is not None:
                (ngi, ngii, hidebit) = ret
                secretFlag = secretFlag + hidebit
                iw[imgRowFlag][imgColFlag] = ngi
                iw[newRowFlag][newColFlag] = ngii
            (imgRowFlag, imgColFlag) = self.zigzagNextIndex(newRowFlag, newColFlag, io.shape[1])
        return iw
    
    def extract(self, stegaImgPath):
        io=cv.imread(stegaImgPath, cv.IMREAD_GRAYSCALE) # 读入载体图像
        iw=np.array(io)
        a, b = io.shape[0], io.shape[1]
        c = 3 if len(io.shape) == 3 else 1
        # total_pixels = a * (b // 2) * c
        total_pixels = 100
        processed_pixels = 0
        imgRowFlag = 0
        imgColFlag = 0
        secretBits = np.array([], dtype=np.uint8)

        isEnding = []
        
        cacheSecretBits = np.array([], dtype=np.uint8)
        global progress_bar_ratio
        while imgRowFlag < io.shape[0]:
            (newRowFlag, newColFlag) =  self.zigzagNextIndex(imgRowFlag, imgColFlag, io.shape[1])
            gi = io[imgRowFlag][imgColFlag]
            gii = io[newRowFlag][newColFlag]
            newSecretBits = self.extractFromPair(gi, gii)
            # print(newSecretBits)
            # prin(cacheSecretBits)
            # print(secretBits)
            if newSecretBits is not None:
                for b in newSecretBits:
                    if b == 0:
                        isEnding.append(0)
                        if len(isEnding) == 32:
                            break
                    else:
                        isEnding = []
                        break
                if len(isEnding) == 0:
                    cacheSecretBits = np.append(cacheSecretBits, newSecretBits)
                    secretBits = np.append(secretBits, cacheSecretBits)
                    cacheSecretBits = np.array([], dtype=np.uint8)
                elif len(isEnding) == 32:
                    hiddenbytes = bytes(np.packbits(secretBits))
                    hiddentxt = hiddenbytes.decode()
                    return hiddentxt
                else:
                    assert(len(isEnding) < 32)
                    cacheSecretBits = np.append(cacheSecretBits, newSecretBits)
                processed_pixels += 1
                progress_bar_ratio=processed_pixels/total_pixels
                progress_bar_ratio=int(round(progress_bar_ratio,2)*100)

            (imgRowFlag, imgColFlag) = self.zigzagNextIndex(newRowFlag, newColFlag, io.shape[1])
        hiddenbytes = bytes(np.packbits(secretBits))
        hiddentxt = hiddenbytes.decode()
        return hiddentxt

def main():
    parser = argparse.ArgumentParser("PVD algorithm")
    parser.add_argument("action", 
            choices = ["hide", "extract"],
            help="Hide or extract secret")
    parser.add_argument("--secretData", 
            type=pathlib.Path,
            default=pathlib.Path("secret.txt"),
            help="The text file that contains the secret")
    parser.add_argument("--coverImage", 
            type=pathlib.Path,
            default = pathlib.Path("lena.bmp"),
            help="The cover image file path")
    parser.add_argument("--stegedImage",
            type=pathlib.Path,
            default = pathlib.Path("hide.bmp"),
            help="The steged image which contains the secret data")

    try:
        args = parser.parse_args()
        rangeBits = [3, 3, 4, 5, 6, 7]
        pvd = PVD(rangeBits)
        if args.action == "hide":
            with open(args.secretData) as f:
                msgRaw = f.read()
            iw = pvd.hide(str(args.coverImage), msgRaw)
            cv.imwrite(str(args.stegedImage), iw)
            cv.imshow("hiden", iw)
            cv.waitKey(0)
        else:
            secretBits = pvd.extract(str(args.stegedImage))
            hiddenbytes = bytes(np.packbits(secretBits))
            hiddentxt = hiddenbytes.decode()
            print(hiddentxt)
    except Exception as e:
        raise e
        print(e)
        sys.exit()
    
if __name__ == "__main__":
    main()
